package com.ai.agent.demo.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 查询扩展器 Demo
 */
@Component
public class MultiQueryExpanderDemo {

    /**
     * 聊天客户端构建器，用于创建与大语言模型交互的客户端
     */
    private final ChatClient.Builder chatClientBuilder;

    /**
     * 构造函数，注入DashScope聊天模型
     * @param dashscopeChatModel DashScope提供的聊天模型实例
     */
    public MultiQueryExpanderDemo(ChatModel dashscopeChatModel) {
        this.chatClientBuilder = ChatClient.builder(dashscopeChatModel);
    }

    /**
     * 将单个查询扩展为多个相关查询
     * @param query 原始查询文本
     * @return 扩展后的查询列表（默认生成3个变体）
     */
    public List<Query> expand(String query) {
        MultiQueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(chatClientBuilder)
                .numberOfQueries(3)
                .build();
        List<Query> queries = queryExpander.expand(new Query("亚健康怎么治疗啊？"));
        return queries;
    }
}
